{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Prediction Verification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This script demonstrates how to train a single model class, embed the model, and solve the optimization problem for *regression* problems (i.e., continuous outcome prediction). We fix a sample from our generated data and solve the optimization problem with all elements of $\\mathbf{x}$ equal to our data. In general, we might have some elements of $\\mathbf{x}$ that are fixed, called our \"conceptual variables,\" and the remaining indices are our decision variables. By fixing all elements of $\\mathbf{x}$, we can verify that the model prediction matches the original sklearn model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the relevant packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "from sklearn.utils.extmath import cartesian\n",
    "import time\n",
    "import sys\n",
    "import os\n",
    "import time\n",
    "\n",
    "from sklearn.metrics import roc_auc_score, r2_score, mean_squared_error\n",
    "from sklearn.cluster import KMeans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import opticl\n",
    "from pyomo import environ\n",
    "from pyomo.environ import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize data\n",
    "We will work with a basic dataset from `sklearn`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>col0</th>\n",
       "      <th>col1</th>\n",
       "      <th>col2</th>\n",
       "      <th>col3</th>\n",
       "      <th>col4</th>\n",
       "      <th>col5</th>\n",
       "      <th>col6</th>\n",
       "      <th>col7</th>\n",
       "      <th>col8</th>\n",
       "      <th>col9</th>\n",
       "      <th>col10</th>\n",
       "      <th>col11</th>\n",
       "      <th>col12</th>\n",
       "      <th>col13</th>\n",
       "      <th>col14</th>\n",
       "      <th>col15</th>\n",
       "      <th>col16</th>\n",
       "      <th>col17</th>\n",
       "      <th>col18</th>\n",
       "      <th>col19</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.584662</td>\n",
       "      <td>0.018317</td>\n",
       "      <td>-0.218921</td>\n",
       "      <td>0.554458</td>\n",
       "      <td>0.532759</td>\n",
       "      <td>-1.507076</td>\n",
       "      <td>-0.004103</td>\n",
       "      <td>1.706662</td>\n",
       "      <td>0.064656</td>\n",
       "      <td>1.135484</td>\n",
       "      <td>1.301202</td>\n",
       "      <td>-0.376086</td>\n",
       "      <td>1.203994</td>\n",
       "      <td>-1.226023</td>\n",
       "      <td>-0.776604</td>\n",
       "      <td>-0.123396</td>\n",
       "      <td>0.958069</td>\n",
       "      <td>0.260800</td>\n",
       "      <td>-1.072332</td>\n",
       "      <td>0.712571</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.127315</td>\n",
       "      <td>0.664080</td>\n",
       "      <td>-0.588780</td>\n",
       "      <td>0.317858</td>\n",
       "      <td>-0.426667</td>\n",
       "      <td>-0.072892</td>\n",
       "      <td>0.068032</td>\n",
       "      <td>-0.628463</td>\n",
       "      <td>0.858286</td>\n",
       "      <td>-0.070950</td>\n",
       "      <td>0.307331</td>\n",
       "      <td>0.186212</td>\n",
       "      <td>1.004093</td>\n",
       "      <td>0.339488</td>\n",
       "      <td>1.394081</td>\n",
       "      <td>-1.311324</td>\n",
       "      <td>0.995704</td>\n",
       "      <td>0.873006</td>\n",
       "      <td>0.757328</td>\n",
       "      <td>0.292931</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-0.097976</td>\n",
       "      <td>0.578025</td>\n",
       "      <td>-0.776809</td>\n",
       "      <td>-1.559095</td>\n",
       "      <td>-0.670203</td>\n",
       "      <td>0.465758</td>\n",
       "      <td>1.130605</td>\n",
       "      <td>0.510315</td>\n",
       "      <td>-0.730213</td>\n",
       "      <td>2.378240</td>\n",
       "      <td>-0.792741</td>\n",
       "      <td>-0.226675</td>\n",
       "      <td>1.131265</td>\n",
       "      <td>-0.253619</td>\n",
       "      <td>1.393016</td>\n",
       "      <td>-0.284635</td>\n",
       "      <td>-2.739142</td>\n",
       "      <td>1.050003</td>\n",
       "      <td>-0.306036</td>\n",
       "      <td>-0.457663</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.667930</td>\n",
       "      <td>-1.083960</td>\n",
       "      <td>-0.633375</td>\n",
       "      <td>0.955877</td>\n",
       "      <td>0.261402</td>\n",
       "      <td>-1.520746</td>\n",
       "      <td>-1.217973</td>\n",
       "      <td>-0.057682</td>\n",
       "      <td>0.852833</td>\n",
       "      <td>0.004939</td>\n",
       "      <td>-0.788276</td>\n",
       "      <td>0.976803</td>\n",
       "      <td>0.579087</td>\n",
       "      <td>1.919169</td>\n",
       "      <td>-0.536705</td>\n",
       "      <td>-2.160269</td>\n",
       "      <td>-1.993137</td>\n",
       "      <td>0.844518</td>\n",
       "      <td>-0.079740</td>\n",
       "      <td>0.477895</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.680070</td>\n",
       "      <td>-0.517094</td>\n",
       "      <td>-0.174703</td>\n",
       "      <td>2.190700</td>\n",
       "      <td>-1.896361</td>\n",
       "      <td>-0.287308</td>\n",
       "      <td>0.901487</td>\n",
       "      <td>-0.248635</td>\n",
       "      <td>0.213534</td>\n",
       "      <td>-0.319802</td>\n",
       "      <td>-0.646917</td>\n",
       "      <td>0.986335</td>\n",
       "      <td>0.248799</td>\n",
       "      <td>0.043669</td>\n",
       "      <td>0.495211</td>\n",
       "      <td>-0.997027</td>\n",
       "      <td>2.528326</td>\n",
       "      <td>-0.296641</td>\n",
       "      <td>1.331457</td>\n",
       "      <td>-0.226314</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>145</th>\n",
       "      <td>0.614726</td>\n",
       "      <td>1.495885</td>\n",
       "      <td>2.293718</td>\n",
       "      <td>-0.214654</td>\n",
       "      <td>1.021248</td>\n",
       "      <td>0.353567</td>\n",
       "      <td>-0.477124</td>\n",
       "      <td>1.037039</td>\n",
       "      <td>-1.019520</td>\n",
       "      <td>-0.348984</td>\n",
       "      <td>0.524750</td>\n",
       "      <td>-0.830011</td>\n",
       "      <td>0.599213</td>\n",
       "      <td>0.672620</td>\n",
       "      <td>0.606404</td>\n",
       "      <td>0.675454</td>\n",
       "      <td>-0.035990</td>\n",
       "      <td>-1.470237</td>\n",
       "      <td>1.005687</td>\n",
       "      <td>2.428877</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>146</th>\n",
       "      <td>-2.506441</td>\n",
       "      <td>-0.135977</td>\n",
       "      <td>0.956122</td>\n",
       "      <td>-0.237942</td>\n",
       "      <td>1.155288</td>\n",
       "      <td>-1.942589</td>\n",
       "      <td>1.122328</td>\n",
       "      <td>-0.106794</td>\n",
       "      <td>1.192686</td>\n",
       "      <td>-2.114164</td>\n",
       "      <td>0.438166</td>\n",
       "      <td>-0.705841</td>\n",
       "      <td>0.282676</td>\n",
       "      <td>1.451429</td>\n",
       "      <td>0.621083</td>\n",
       "      <td>-0.797270</td>\n",
       "      <td>-0.997020</td>\n",
       "      <td>-0.826097</td>\n",
       "      <td>-2.037201</td>\n",
       "      <td>-0.618037</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>147</th>\n",
       "      <td>-1.004765</td>\n",
       "      <td>-0.520928</td>\n",
       "      <td>-0.068925</td>\n",
       "      <td>-0.498428</td>\n",
       "      <td>-0.304744</td>\n",
       "      <td>1.707986</td>\n",
       "      <td>-0.486471</td>\n",
       "      <td>-0.491599</td>\n",
       "      <td>-0.215651</td>\n",
       "      <td>0.754817</td>\n",
       "      <td>-0.162497</td>\n",
       "      <td>0.752761</td>\n",
       "      <td>1.254887</td>\n",
       "      <td>-0.454928</td>\n",
       "      <td>-0.091675</td>\n",
       "      <td>-0.859160</td>\n",
       "      <td>0.088282</td>\n",
       "      <td>0.241109</td>\n",
       "      <td>0.007260</td>\n",
       "      <td>1.986539</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>148</th>\n",
       "      <td>-0.840332</td>\n",
       "      <td>-0.880045</td>\n",
       "      <td>-0.762087</td>\n",
       "      <td>0.475760</td>\n",
       "      <td>0.373943</td>\n",
       "      <td>-1.642020</td>\n",
       "      <td>-0.638848</td>\n",
       "      <td>1.955754</td>\n",
       "      <td>-0.009185</td>\n",
       "      <td>1.271048</td>\n",
       "      <td>-0.879264</td>\n",
       "      <td>1.053166</td>\n",
       "      <td>-0.022013</td>\n",
       "      <td>0.473292</td>\n",
       "      <td>-0.288755</td>\n",
       "      <td>-0.653662</td>\n",
       "      <td>0.855587</td>\n",
       "      <td>0.338300</td>\n",
       "      <td>1.228505</td>\n",
       "      <td>-0.675083</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>149</th>\n",
       "      <td>0.527951</td>\n",
       "      <td>-0.397126</td>\n",
       "      <td>0.573265</td>\n",
       "      <td>0.309318</td>\n",
       "      <td>2.917309</td>\n",
       "      <td>-0.563986</td>\n",
       "      <td>1.153213</td>\n",
       "      <td>0.079840</td>\n",
       "      <td>0.937555</td>\n",
       "      <td>-1.561120</td>\n",
       "      <td>1.098689</td>\n",
       "      <td>-0.108898</td>\n",
       "      <td>0.318972</td>\n",
       "      <td>1.312895</td>\n",
       "      <td>1.040038</td>\n",
       "      <td>0.886090</td>\n",
       "      <td>1.290993</td>\n",
       "      <td>0.264868</td>\n",
       "      <td>-0.831173</td>\n",
       "      <td>0.023357</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>150 rows × 20 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         col0      col1      col2      col3      col4      col5      col6  \\\n",
       "0    1.584662  0.018317 -0.218921  0.554458  0.532759 -1.507076 -0.004103   \n",
       "1    0.127315  0.664080 -0.588780  0.317858 -0.426667 -0.072892  0.068032   \n",
       "2   -0.097976  0.578025 -0.776809 -1.559095 -0.670203  0.465758  1.130605   \n",
       "3   -0.667930 -1.083960 -0.633375  0.955877  0.261402 -1.520746 -1.217973   \n",
       "4    0.680070 -0.517094 -0.174703  2.190700 -1.896361 -0.287308  0.901487   \n",
       "..        ...       ...       ...       ...       ...       ...       ...   \n",
       "145  0.614726  1.495885  2.293718 -0.214654  1.021248  0.353567 -0.477124   \n",
       "146 -2.506441 -0.135977  0.956122 -0.237942  1.155288 -1.942589  1.122328   \n",
       "147 -1.004765 -0.520928 -0.068925 -0.498428 -0.304744  1.707986 -0.486471   \n",
       "148 -0.840332 -0.880045 -0.762087  0.475760  0.373943 -1.642020 -0.638848   \n",
       "149  0.527951 -0.397126  0.573265  0.309318  2.917309 -0.563986  1.153213   \n",
       "\n",
       "         col7      col8      col9     col10     col11     col12     col13  \\\n",
       "0    1.706662  0.064656  1.135484  1.301202 -0.376086  1.203994 -1.226023   \n",
       "1   -0.628463  0.858286 -0.070950  0.307331  0.186212  1.004093  0.339488   \n",
       "2    0.510315 -0.730213  2.378240 -0.792741 -0.226675  1.131265 -0.253619   \n",
       "3   -0.057682  0.852833  0.004939 -0.788276  0.976803  0.579087  1.919169   \n",
       "4   -0.248635  0.213534 -0.319802 -0.646917  0.986335  0.248799  0.043669   \n",
       "..        ...       ...       ...       ...       ...       ...       ...   \n",
       "145  1.037039 -1.019520 -0.348984  0.524750 -0.830011  0.599213  0.672620   \n",
       "146 -0.106794  1.192686 -2.114164  0.438166 -0.705841  0.282676  1.451429   \n",
       "147 -0.491599 -0.215651  0.754817 -0.162497  0.752761  1.254887 -0.454928   \n",
       "148  1.955754 -0.009185  1.271048 -0.879264  1.053166 -0.022013  0.473292   \n",
       "149  0.079840  0.937555 -1.561120  1.098689 -0.108898  0.318972  1.312895   \n",
       "\n",
       "        col14     col15     col16     col17     col18     col19  \n",
       "0   -0.776604 -0.123396  0.958069  0.260800 -1.072332  0.712571  \n",
       "1    1.394081 -1.311324  0.995704  0.873006  0.757328  0.292931  \n",
       "2    1.393016 -0.284635 -2.739142  1.050003 -0.306036 -0.457663  \n",
       "3   -0.536705 -2.160269 -1.993137  0.844518 -0.079740  0.477895  \n",
       "4    0.495211 -0.997027  2.528326 -0.296641  1.331457 -0.226314  \n",
       "..        ...       ...       ...       ...       ...       ...  \n",
       "145  0.606404  0.675454 -0.035990 -1.470237  1.005687  2.428877  \n",
       "146  0.621083 -0.797270 -0.997020 -0.826097 -2.037201 -0.618037  \n",
       "147 -0.091675 -0.859160  0.088282  0.241109  0.007260  1.986539  \n",
       "148 -0.288755 -0.653662  0.855587  0.338300  1.228505 -0.675083  \n",
       "149  1.040038  0.886090  1.290993  0.264868 -0.831173  0.023357  \n",
       "\n",
       "[150 rows x 20 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.datasets import make_regression\n",
    "from sklearn.model_selection import train_test_split\n",
    "X, y = make_regression(n_samples=200, n_features = 20, random_state=1)\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y,\n",
    "                                                    random_state=1)\n",
    "X_train = pd.DataFrame(X_train).add_prefix('col')\n",
    "X_test = pd.DataFrame(X_test).add_prefix('col')\n",
    "X_train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the chosen model type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# alg = 'rf' \n",
    "alg = 'gbm'\n",
    "task_type = 'continuous'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The user can optionally select a manual parameter grid for the cross-validation procedure. We implement a default parameter grid; see **run_MLmodels.py** for details on the tuned parameters. If you wish to use the default, leave ```parameter_grid = None``` (or do not specify any grid)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_grid = None\n",
    "# parameter_grid = {'hidden_layer_sizes': [(5,),(10,)]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------- Initialize grid  ----------------\n",
      "------------- Running model  ----------------\n",
      "Algorithm = gbm, metric = None\n",
      "saving... results/gbm_temp_trained.pkl\n",
      "------------- Model evaluation  ----------------\n",
      "-------------------training evaluation-----------------------\n",
      "Train MSE: 4314.00082576947\n",
      "Train R2: 0.8939305706172604\n",
      "-------------------testing evaluation-----------------------\n",
      "Test MSE: 17814.940522763252\n",
      "Test R2: 0.6170544988313675\n"
     ]
    }
   ],
   "source": [
    "s = 1\n",
    "version = 'test'\n",
    "outcome = 'temp'\n",
    "\n",
    "model_save = 'results/%s/%s_%s_model.csv' % (alg, version, outcome)\n",
    "\n",
    "alg_run = alg if alg != 'rf' else 'rf_shallow'\n",
    "m, perf = opticl.run_model(X_train, y_train, X_test, y_test, alg_run, outcome, task = task_type, \n",
    "                       seed = s, cv_folds = 5, \n",
    "                       # The user can manually specify the parameter grid for cross-validation if desired\n",
    "                       parameter_grid = parameter_grid,\n",
    "                       save_path = model_save,\n",
    "                       save = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After training the model, we will save the trained model in the format needed for embedding the constraints. See **constraint_learning.py** for the specific format that is extracted per method. We also save the performance of the model to use in the automated model selection pipeline (if desired).\n",
    "\n",
    "We also create the save directory if it does not exist.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('results/%s/' % alg):\n",
    "    os.makedirs('results/%s/' % alg)\n",
    "    \n",
    "constraintL = opticl.ConstraintLearning(X_train, y_train, m, alg)\n",
    "constraint_add = constraintL.constraint_extrapolation(task_type)\n",
    "constraint_add.to_csv(model_save, index = False)\n",
    "\n",
    "perf.to_csv('results/%s/%s_%s_performance.csv' % (alg, version, outcome), index= False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Tree_id</th>\n",
       "      <th>ID</th>\n",
       "      <th>col0</th>\n",
       "      <th>col1</th>\n",
       "      <th>col2</th>\n",
       "      <th>col3</th>\n",
       "      <th>col4</th>\n",
       "      <th>col5</th>\n",
       "      <th>col6</th>\n",
       "      <th>col7</th>\n",
       "      <th>...</th>\n",
       "      <th>col14</th>\n",
       "      <th>col15</th>\n",
       "      <th>col16</th>\n",
       "      <th>col17</th>\n",
       "      <th>col18</th>\n",
       "      <th>col19</th>\n",
       "      <th>threshold</th>\n",
       "      <th>prediction</th>\n",
       "      <th>initial_prediction</th>\n",
       "      <th>learning_rate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.362847</td>\n",
       "      <td>-147.677247</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.040611</td>\n",
       "      <td>-147.677247</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.362847</td>\n",
       "      <td>26.265507</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-0.040612</td>\n",
       "      <td>26.265507</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-0.362848</td>\n",
       "      <td>-36.183599</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>19</td>\n",
       "      <td>2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-0.653505</td>\n",
       "      <td>29.789185</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>19</td>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1.771010</td>\n",
       "      <td>117.421701</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>19</td>\n",
       "      <td>3</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.913692</td>\n",
       "      <td>117.421701</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>19</td>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1.771010</td>\n",
       "      <td>71.336522</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>19</td>\n",
       "      <td>4</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>-1.913693</td>\n",
       "      <td>71.336522</td>\n",
       "      <td>3.181564</td>\n",
       "      <td>0.2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>160 rows × 26 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   Tree_id ID  col0 col1  col2 col3  col4  col5  col6  col7  ...  col14  \\\n",
       "0        0  1   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0        0  1   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0        0  2   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0        0  2   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0        0  3   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "..     ... ..   ...  ...   ...  ...   ...   ...   ...   ...  ...    ...   \n",
       "0       19  2   0.0    0   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0       19  3   0.0   -1   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0       19  3   0.0    1   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0       19  4   0.0   -1   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "0       19  4   0.0   -1   0.0    0   0.0   0.0   0.0   0.0  ...      0   \n",
       "\n",
       "    col15  col16 col17 col18  col19 threshold  prediction  initial_prediction  \\\n",
       "0     0.0    0.0     0   0.0    0.0  0.362847 -147.677247            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0  0.040611 -147.677247            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0  0.362847   26.265507            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0 -0.040612   26.265507            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0 -0.362848  -36.183599            3.181564   \n",
       "..    ...    ...   ...   ...    ...       ...         ...                 ...   \n",
       "0     0.0    0.0     0   0.0    0.0 -0.653505   29.789185            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0 -1.771010  117.421701            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0  1.913692  117.421701            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0 -1.771010   71.336522            3.181564   \n",
       "0     0.0    0.0     0   0.0    0.0 -1.913693   71.336522            3.181564   \n",
       "\n",
       "   learning_rate  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "..           ...  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "0            0.2  \n",
       "\n",
       "[160 rows x 26 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "constraint_add"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check: what should the result be for our sample observation, if all x are fixed?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Choose sample to test\n",
    "This will be the observation (\"patient\") that we feed into the optimization model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_id = 1\n",
    "sample = X_train.loc[sample_id:sample_id,:].reset_index(drop = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate model prediction directly in sklearn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([182.75931286])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.predict(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization formulation\n",
    "We will embed the model trained above. The model could also be selected using the model selection pipeline, which we demonstrate in the WFP example script.\n",
    "\n",
    "If manually specifying the model, as we are here, the key elements of the ``model_master`` dataframe are:\n",
    "- model_type: algorithm name.\n",
    "- outcome: name of outcome of interest; this is relevant in the case of multiple learned outcomes.\n",
    "- save_path: file name of the extracted model.\n",
    "- objective: the weight of the objective if it should be included as an additive term in the objective. A weight of 0 omits it from the objective entirely.\n",
    "- lb/ub: the lower (or upper) bound that we wish to apply to the learned outcome. If there is no bound, it should be set to ``None``.\n",
    "\n",
    "In this case, we set the outcome to be our only objective term, which will allow us to verify that the predictions are consistent between the embedded model and the sklearn prediction function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_master = pd.DataFrame(columns = ['model_type','outcome','save_path','lb','ub','objective'])\n",
    "\n",
    "model_master.loc[0,'model_type'] = alg\n",
    "model_master.loc[0,'save_path'] = 'results/%s/%s_%s_model.csv' % (alg, version, outcome)\n",
    "model_master.loc[0,'outcome'] = outcome\n",
    "model_master.loc[0,'objective'] = 1\n",
    "model_master.loc[0,'ub'] = None\n",
    "model_master.loc[0,'lb'] = None\n",
    "model_master.loc[0,'task'] = task_type\n",
    "model_master['SCM_counterfactuals'] = None\n",
    "model_master['features'] = [[col for col in X_train.columns]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solve with Pyomo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model_pyo = ConcreteModel()\n",
    "\n",
    "## We will create our x decision variables, and fix them all to our sample's values for model verification.\n",
    "N = X_train.columns\n",
    "model_pyo.x = Var(N, domain=Reals)\n",
    "\n",
    "def fix_value(model_pyo, index):\n",
    "    return model_pyo.x[index] == sample.loc[0,index]\n",
    "\n",
    "model_pyo.Constraint1 = Constraint(N, rule=fix_value)\n",
    "\n",
    "## Specify any non-learned objective components - none here \n",
    "model_pyo.OBJ = Objective(expr=0, sense=minimize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding objective function for temp\n"
     ]
    }
   ],
   "source": [
    "final_model_pyo = opticl.optimization_MIP(model_pyo, model_pyo.x, model_master, X_train, tr = False)\n",
    "# final_model_pyo.pprint()\n",
    "opt = SolverFactory('gurobi')\n",
    "results = opt.solve(final_model_pyo) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check for equality between sklearn and embedded models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True outcome: 182.759\n",
      "Pyomo output: 182.759\n"
     ]
    }
   ],
   "source": [
    "print(\"True outcome: %.3f\" % m.predict(sample)[0])\n",
    "print(\"Pyomo output: %.3f\" % final_model_pyo.OBJ())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
